#!/bin/bash

# =============================================================================
# Dify 生产环境维护脚本
# =============================================================================

set -euo pipefail
IFS=$'\n\t'

# --- 配置 ---
SCRIPT_VERSION="1.0.0"
BASE_DIR="/opt/dify"
COMPOSE_FILE="$BASE_DIR/docker-compose.production.yml"
ENV_FILE="$BASE_DIR/config/.env.production"
BACKUP_DIR="$BASE_DIR/backups"
LOG_DIR="$BASE_DIR/logs"

# --- 颜色输出 ---
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# --- 日志函数 ---
log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
log_warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# --- 检查函数 ---
check_environment() {
    if [[ ! -d "$BASE_DIR" ]]; then
        log_error "Dify 目录不存在: $BASE_DIR"
        exit 1
    fi
    
    if [[ ! -f "$COMPOSE_FILE" ]]; then
        log_error "Docker Compose 文件不存在: $COMPOSE_FILE"
        exit 1
    fi
    
    cd "$BASE_DIR"
}

# --- 服务管理函数 ---
start_services() {
    log_info "启动 Dify 服务..."
    
    check_environment
    docker compose -f docker-compose.production.yml up -d
    
    log_success "Dify 服务启动完成"
    sleep 5
    status_services
}

stop_services() {
    log_info "停止 Dify 服务..."
    
    check_environment
    docker compose -f docker-compose.production.yml down
    
    log_success "Dify 服务停止完成"
}

restart_services() {
    log_info "重启 Dify 服务..."
    
    check_environment
    docker compose -f docker-compose.production.yml restart
    
    log_success "Dify 服务重启完成"
    sleep 5
    status_services
}

status_services() {
    log_info "检查 Dify 服务状态..."
    
    check_environment
    echo ""
    echo "=== Docker 服务状态 ==="
    docker compose -f docker-compose.production.yml ps
    
    echo ""
    echo "=== 服务健康检查 ==="
    local healthy=0
    local total=0
    
    for service in postgres redis weaviate sandbox api worker web nginx; do
        total=$((total + 1))
        if docker compose -f docker-compose.production.yml ps --format json | jq -r ".[] | select(.Service == \"$service\") | .Health" | grep -q "healthy"; then
            echo "✅ $service: 健康"
            healthy=$((healthy + 1))
        else
            echo "❌ $service: 不健康"
        fi
    done
    
    echo ""
    echo "健康状态: $healthy/$total 服务正常"
}

logs_services() {
    check_environment
    
    if [[ $# -eq 0 ]]; then
        log_info "显示所有服务日志..."
        docker compose -f docker-compose.production.yml logs -f --tail=100
    else
        log_info "显示 $1 服务日志..."
        docker compose -f docker-compose.production.yml logs -f --tail=100 "$1"
    fi
}

# --- 备份和恢复函数 ---
backup_data() {
    log_info "开始备份 Dify 数据..."
    
    check_environment
    
    local date=$(date +%Y%m%d_%H%M%S)
    local backup_file="$BACKUP_DIR/dify_backup_$date.tar.gz"
    
    # 创建备份目录
    mkdir -p "$BACKUP_DIR"
    
    # 停止服务
    log_info "停止服务以确保数据一致性..."
    docker compose -f docker-compose.production.yml stop
    
    # 创建备份
    log_info "创建数据备份..."
    tar -czf "$backup_file" \
        --exclude="$BACKUP_DIR" \
        --exclude="$LOG_DIR" \
        --exclude="*.log" \
        -C "$BASE_DIR" \
        data/ config/
    
    # 重启服务
    log_info "重新启动服务..."
    docker compose -f docker-compose.production.yml start
    
    log_success "备份完成: $backup_file"
    
    # 清理旧备份（保留7天）
    find "$BACKUP_DIR" -name "dify_backup_*.tar.gz" -mtime +7 -delete
    log_info "已清理7天前的旧备份"
}

restore_data() {
    if [[ $# -eq 0 ]]; then
        log_error "请指定备份文件路径"
        echo "用法: $0 restore <backup_file>"
        exit 1
    fi
    
    local backup_file="$1"
    
    if [[ ! -f "$backup_file" ]]; then
        log_error "备份文件不存在: $backup_file"
        exit 1
    fi
    
    log_warning "这将恢复数据并覆盖当前数据，是否继续? (y/N)"
    read -r confirm
    if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
        log_info "操作已取消"
        exit 0
    fi
    
    check_environment
    
    # 停止服务
    log_info "停止服务..."
    docker compose -f docker-compose.production.yml down
    
    # 备份当前数据
    local current_backup="$BACKUP_DIR/before_restore_$(date +%Y%m%d_%H%M%S).tar.gz"
    log_info "备份当前数据到: $current_backup"
    tar -czf "$current_backup" -C "$BASE_DIR" data/ config/
    
    # 恢复数据
    log_info "恢复数据..."
    tar -xzf "$backup_file" -C "$BASE_DIR"
    
    # 启动服务
    log_info "启动服务..."
    docker compose -f docker-compose.production.yml up -d
    
    log_success "数据恢复完成"
}

# --- 更新函数 ---
update_dify() {
    log_info "开始更新 Dify..."
    
    check_environment
    
    # 备份数据
    backup_data
    
    # 拉取新镜像
    log_info "拉取最新镜像..."
    docker compose -f docker-compose.production.yml pull
    
    # 重启服务
    log_info "重启服务以应用更新..."
    docker compose -f docker-compose.production.yml up -d
    
    # 清理旧镜像
    log_info "清理未使用的镜像..."
    docker image prune -f
    
    log_success "Dify 更新完成"
    sleep 10
    status_services
}

# --- 系统监控函数 ---
system_info() {
    log_info "系统资源使用情况:"
    
    echo ""
    echo "=== CPU 使用率 ==="
    top -bn1 | grep "Cpu(s)" | awk '{print $2}' | awk -F'%' '{print "CPU: "$1"%"}'
    
    echo ""
    echo "=== 内存使用 ==="
    free -h | grep "Mem:" | awk '{printf "内存: %s / %s (%.1f%%)\n", $3, $2, ($3/$2)*100}'
    
    echo ""
    echo "=== 磁盘使用 ==="
    df -h / | tail -1 | awk '{printf "磁盘: %s / %s (%s)\n", $3, $2, $5}'
    
    echo ""
    echo "=== Docker 资源使用 ==="
    docker stats --no-stream --format "table {{.Container}}\t{{.CPUPerc}}\t{{.MemUsage}}\t{{.NetIO}}"
    
    echo ""
    echo "=== 网络连接 ==="
    ss -tuln | grep -E "(80|443|5432|6379|8080)" | wc -l | awk '{print "活动连接数: "$1}'
}

# --- 清理函数 ---
cleanup_system() {
    log_info "清理系统资源..."
    
    # 清理 Docker
    log_info "清理 Docker 资源..."
    docker system prune -f
    docker volume prune -f
    
    # 清理日志
    log_info "清理旧日志..."
    find "$LOG_DIR" -name "*.log" -mtime +30 -delete 2>/dev/null || true
    
    # 清理临时文件
    log_info "清理临时文件..."
    find /tmp -name "nginx_cache" -type d -exec rm -rf {} + 2>/dev/null || true
    
    log_success "系统清理完成"
}

# --- SSL 证书管理 ---
renew_ssl() {
    log_info "更新 SSL 证书..."
    
    if command -v certbot &> /dev/null; then
        certbot renew --nginx --quiet
        
        # 重启 nginx 以加载新证书
        docker compose -f docker-compose.production.yml restart nginx
        
        log_success "SSL 证书更新完成"
    else
        log_warning "certbot 未安装，跳过 SSL 证书更新"
    fi
}

# --- 配置管理 ---
edit_config() {
    local config_file="$ENV_FILE"
    
    if [[ ! -f "$config_file" ]]; then
        log_error "配置文件不存在: $config_file"
        exit 1
    fi
    
    log_info "编辑配置文件: $config_file"
    
    # 使用默认编辑器
    ${EDITOR:-nano} "$config_file"
    
    log_warning "配置已修改，需要重启服务以生效"
    read -p "是否现在重启服务? (y/N): " -r confirm
    if [[ "$confirm" =~ ^[Yy]$ ]]; then
        restart_services
    fi
}

# --- 帮助函数 ---
show_help() {
    echo "Dify 生产环境维护脚本 v$SCRIPT_VERSION"
    echo ""
    echo "用法: $0 <command> [options]"
    echo ""
    echo "命令:"
    echo "  start           启动 Dify 服务"
    echo "  stop            停止 Dify 服务"
    echo "  restart         重启 Dify 服务"
    echo "  status          查看服务状态"
    echo "  logs [service]  查看日志 (可选指定服务名)"
    echo ""
    echo "  backup          备份数据"
    echo "  restore <file>  恢复数据"
    echo "  update          更新 Dify"
    echo ""
    echo "  info            显示系统信息"
    echo "  cleanup         清理系统资源"
    echo "  ssl-renew       更新 SSL 证书"
    echo "  config          编辑配置文件"
    echo ""
    echo "  help            显示此帮助"
    echo ""
    echo "示例:"
    echo "  $0 start                    # 启动服务"
    echo "  $0 logs api                 # 查看 API 服务日志"
    echo "  $0 backup                   # 备份数据"
    echo "  $0 restore backup.tar.gz    # 恢复数据"
}

# --- 主函数 ---
main() {
    case "${1:-}" in
        start)
            start_services
            ;;
        stop)
            stop_services
            ;;
        restart)
            restart_services
            ;;
        status)
            status_services
            ;;
        logs)
            logs_services "${2:-}"
            ;;
        backup)
            backup_data
            ;;
        restore)
            restore_data "${2:-}"
            ;;
        update)
            update_dify
            ;;
        info)
            system_info
            ;;
        cleanup)
            cleanup_system
            ;;
        ssl-renew)
            renew_ssl
            ;;
        config)
            edit_config
            ;;
        help|--help|-h)
            show_help
            ;;
        *)
            log_error "未知命令: ${1:-}"
            echo ""
            show_help
            exit 1
            ;;
    esac
}

# 检查依赖
if ! command -v docker &> /dev/null; then
    log_error "Docker 未安装"
    exit 1
fi

if ! command -v jq &> /dev/null; then
    log_warning "建议安装 jq 以获得更好的状态显示"
fi

# 运行主函数
main "$@" 